/*
 * wifi_sta.c
 * 
 * 使用wpa_supplicant管理WiFi客户端(STA)模式的实现文件
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <signal.h>
#include <errno.h>
#include "wifi_sta.h"

/**
 * 初始化WiFi STA配置为默认值
 */
wifi_ap_error_t wifi_sta_init_config(wifi_sta_config_t *config)
{
    if (config == NULL) {
        return WIFI_AP_ERROR_GENERAL;
    }

    memset(config, 0, sizeof(wifi_sta_config_t));

    // 设置默认值
    strcpy(config->interface, "wlan0");
    strcpy(config->ssid, "");
    strcpy(config->password, "");
    strcpy(config->country_code, "CN");
    strcpy(config->wpa_supplicant_path, "/usr/sbin/wpa_supplicant");
    strcpy(config->config_path, "/tmp/wpa_supplicant.conf");

    return WIFI_AP_SUCCESS;
}

/**
 * 创建wpa_supplicant配置文件
 */
wifi_ap_error_t wifi_sta_create_config_file(const wifi_sta_config_t *config)
{
    FILE *fp;

    if (config == NULL) {
        return WIFI_AP_ERROR_CONFIG;
    }

    // 检查wpa_supplicant是否存在
    if (!file_exists(config->wpa_supplicant_path)) {
        fprintf(stderr, "wpa_supplicant not found at %s\n", config->wpa_supplicant_path);
        return WIFI_AP_ERROR_HOSTAPD_NOT_FOUND;
    }

    // 检查SSID是否为空
    if (strlen(config->ssid) == 0) {
        fprintf(stderr, "SSID cannot be empty\n");
        return WIFI_AP_ERROR_CONFIG;
    }

    // 创建配置文件
    fp = fopen(config->config_path, "w");
    if (fp == NULL) {
        perror("Failed to create wpa_supplicant config file");
        return WIFI_AP_ERROR_CONFIG;
    }

    // 写入配置
    fprintf(fp, "country=%s\n", config->country_code);
    fprintf(fp, "ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev\n");
    fprintf(fp, "update_config=1\n");
    fprintf(fp, "\n");
    fprintf(fp, "network={\n");
    fprintf(fp, "\tssid=\"%s\"\n", config->ssid);
    
    // 根据密码长度决定是否使用加密
    if (strlen(config->password) > 0) {
        fprintf(fp, "\tpsk=\"%s\"\n", config->password);
        fprintf(fp, "\tkey_mgmt=WPA-PSK\n");
    } else {
        fprintf(fp, "\tkey_mgmt=NONE\n");
    }
    
    fprintf(fp, "}\n");

    fclose(fp);
    return WIFI_AP_SUCCESS;
}

/**
 * 启动WiFi STA模式并连接到指定WiFi
 */
wifi_ap_error_t wifi_sta_start(const wifi_sta_config_t *config, wifi_sta_state_t *state)
{
    int pid;
    wifi_ap_error_t ret;
    char cmd[512];

    if (config == NULL || state == NULL) {
        return WIFI_AP_ERROR_GENERAL;
    }

    // 检查是否已经在运行
    if (state->connected) {
        return WIFI_AP_ERROR_ALREADY_RUNNING;
    }

    // 启用WiFi接口
    sprintf(cmd, "ip link set %s up", config->interface);
    ret = execute_command(cmd, "Failed to enable WiFi interface");
    if (ret != WIFI_AP_SUCCESS) {
        return ret;
    }

    // 创建配置文件
    ret = wifi_sta_create_config_file(config);
    if (ret != WIFI_AP_SUCCESS) {
        return ret;
    }

    // 使用fork创建子进程来运行wpa_supplicant
    pid = fork();
    if (pid == -1) {
        perror("Failed to fork process");
        return WIFI_AP_ERROR_PROCESS;
    } else if (pid == 0) {
        // 子进程：执行wpa_supplicant
        execl(config->wpa_supplicant_path, config->wpa_supplicant_path, 
              "-B", "-i", config->interface, "-c", config->config_path, 
              (char *)NULL);
        
        // 如果execl返回，则表示失败
        perror("Failed to execute wpa_supplicant");
        exit(EXIT_FAILURE);
    }

    // 父进程：等待一段时间以确保wpa_supplicant成功启动
    sleep(3);

    // 检查进程是否在运行
    if (kill(pid, 0) == -1) {
        fprintf(stderr, "wpa_supplicant process failed to start\n");
        return WIFI_AP_ERROR_PROCESS;
    }

    // 启动dhclient获取IP地址
    sprintf(cmd, "dhclient %s", config->interface);
    ret = execute_command(cmd, "Warning: Failed to start dhclient");
    // 即使dhclient启动失败，也继续执行，因为我们可以稍后检查IP地址

    // 初始化状态
    state->connected = 1;
    state->pid = pid;
    memset(state->ip_address, 0, sizeof(state->ip_address));
    memcpy(&state->config, config, sizeof(wifi_sta_config_t));

    // 尝试获取IP地址
    wifi_sta_get_ip_address(config->interface, state->ip_address, sizeof(state->ip_address));

    return WIFI_AP_SUCCESS;
}

/**
 * 停止WiFi STA模式
 */
wifi_ap_error_t wifi_sta_stop(wifi_sta_state_t *state)
{
    int ret;
    char cmd[256];

    if (state == NULL || !state->connected) {
        return WIFI_AP_SUCCESS; // 已经停止或无效状态
    }

    // 停止dhclient
    sprintf(cmd, "dhclient -r %s", state->config.interface);
    execute_command(cmd, "Warning: Failed to stop dhclient");

    // 停止wpa_supplicant
    if (kill(state->pid, SIGTERM) == -1) {
        if (errno == ESRCH) {
            // 进程不存在
            state->connected = 0;
            return WIFI_AP_SUCCESS;
        }
        perror("Failed to send SIGTERM to wpa_supplicant");
        return WIFI_AP_ERROR_SIGNAL;
    }

    // 等待进程终止
    ret = waitpid(state->pid, NULL, WNOHANG);
    if (ret == 0) {
        // 进程没有立即终止，等待一段时间
        sleep(2);
        // 如果仍未终止，强制终止
        if (kill(state->pid, SIGKILL) == 0) {
            waitpid(state->pid, NULL, 0);
        }
    }

    // 清理状态
    state->connected = 0;
    state->pid = 0;
    memset(state->ip_address, 0, sizeof(state->ip_address));
    memset(&state->config, 0, sizeof(wifi_sta_config_t));

    return WIFI_AP_SUCCESS;
}

/**
 * 检查WiFi STA连接状态
 */
wifi_ap_error_t wifi_sta_check_status(wifi_sta_state_t *state)
{
    if (state == NULL) {
        return WIFI_AP_ERROR_GENERAL;
    }

    // 检查进程是否在运行
    if (state->connected) {
        if (kill(state->pid, 0) == -1) {
            if (errno == ESRCH) {
                // 进程不存在，但状态显示连接中，更新状态
                state->connected = 0;
                state->pid = 0;
                memset(state->ip_address, 0, sizeof(state->ip_address));
                return WIFI_AP_ERROR_PROCESS;
            }
            perror("Failed to check wpa_supplicant process");
            return WIFI_AP_ERROR_GENERAL;
        }

        // 更新IP地址
        wifi_sta_get_ip_address(state->config.interface, state->ip_address, sizeof(state->ip_address));
    }

    return WIFI_AP_SUCCESS;
}

/**
 * 获取WiFi接口的IP地址
 */
wifi_ap_error_t wifi_sta_get_ip_address(const char *interface, char *ip_address, size_t buf_size)
{
    FILE *fp;
    char cmd[256];
    char line[256];
    char *p;

    if (interface == NULL || ip_address == NULL || buf_size == 0) {
        return WIFI_AP_ERROR_GENERAL;
    }

    // 使用ip命令获取接口IP地址
    sprintf(cmd, "ip addr show %s | grep 'inet ' | awk '{print $2}' | cut -d/ -f1", interface);
    fp = popen(cmd, "r");
    if (fp == NULL) {
        perror("Failed to execute ip command");
        return WIFI_AP_ERROR_PROCESS;
    }

    // 读取命令输出
    if (fgets(line, sizeof(line), fp) != NULL) {
        // 去除换行符
        p = strchr(line, '\n');
        if (p != NULL) {
            *p = '\0';
        }
        
        // 复制IP地址到缓冲区
        strncpy(ip_address, line, buf_size - 1);
        ip_address[buf_size - 1] = '\0';
    } else {
        // 没有获取到IP地址
        ip_address[0] = '\0';
    }

    pclose(fp);
    return WIFI_AP_SUCCESS;
}

/**
 * 清理临时文件和资源
 */
wifi_ap_error_t wifi_sta_cleanup(const wifi_sta_config_t *config)
{
    if (config == NULL) {
        return WIFI_AP_ERROR_GENERAL;
    }

    // 删除配置文件
    if (file_exists(config->config_path)) {
        if (unlink(config->config_path) == -1) {
            perror("Failed to delete wpa_supplicant config file");
            return WIFI_AP_ERROR_GENERAL;
        }
    }

    return WIFI_AP_SUCCESS;
}